
import pandas as pd
import numpy as np
pd.set_option('expand_frame_repr', False)  # 当列太多时不换行
pd.set_option('display.max_rows', 30)

# okex交割合约（usdt本位）资金曲线(原始逻辑)
def equity_curve_for_OKEx_USDT_future_next_open_ori(df, slippage=1 / 1000, c_rate=5 / 10000, leverage_rate=3, face_value=0.01,
                                                min_margin_ratio=1/100):
    # =====下根k线开盘价
    df['next_open'] = df['open'].shift(-1)  # 下根K线的开盘价
    df['next_open'].fillna(value=df['close'], inplace=True)

    # =====找出开仓、平仓的k线
    condition1 = df['pos'] != 0  # 当前周期不为空仓
    condition2 = df['pos'] != df['pos'].shift(1)  # 当前周期和上个周期持仓方向不一样。
    open_pos_condition = condition1 & condition2

    condition1 = df['pos'] != 0  # 当前周期不为空仓
    condition2 = df['pos'] != df['pos'].shift(-1)  # 当前周期和下个周期持仓方向不一样。
    close_pos_condition = condition1 & condition2

    # =====对每次交易进行分组
    df.loc[open_pos_condition, 'start_time'] = df['candle_begin_time']
    df['start_time'].fillna(method='ffill', inplace=True)
    df.loc[df['pos'] == 0, 'start_time'] = pd.NaT

    # =====开始计算资金曲线
    initial_cash = 10000  # 初始资金，默认为10000元
    # ===在开仓时
    # 在open_pos_condition的K线，以开盘价计算买入合约的数量。（当资金量大的时候，可以用5分钟均价）
    df.loc[open_pos_condition, 'contract_num'] = initial_cash * leverage_rate / (face_value * df['open'])
    df['contract_num'] = np.floor(df['contract_num'])  # 对合约张数向下取整
    # 开仓价格：理论开盘价加上相应滑点
    df.loc[open_pos_condition, 'open_pos_price'] = df['open'] * (1 + slippage * df['pos'])
    # 开仓之后剩余的钱，扣除手续费
    df['cash'] = initial_cash - df['open_pos_price'] * face_value * df['contract_num'] * c_rate  # 即保证金

    # ===开仓之后每根K线结束时
    # 买入之后cash，contract_num，open_pos_price不再发生变动
    for _ in ['contract_num', 'open_pos_price', 'cash']:
        df[_].fillna(method='ffill', inplace=True)
    df.loc[df['pos'] == 0, ['contract_num', 'open_pos_price', 'cash']] = None

    # ===在平仓时
    # 平仓价格
    df.loc[close_pos_condition, 'close_pos_price'] = df['next_open'] * (1 - slippage * df['pos'])
    # 平仓之后剩余的钱，扣除手续费
    df.loc[close_pos_condition, 'close_pos_fee'] = df['close_pos_price'] * face_value * df['contract_num'] * c_rate

    # ===计算利润
    # 开仓至今持仓盈亏
    df['profit'] = face_value * df['contract_num'] * (df['close'] - df['open_pos_price']) * df['pos']

    # 平仓时理论额外处理
    df.loc[close_pos_condition, 'profit'] = face_value * df['contract_num'] * (
            df['close_pos_price'] - df['open_pos_price']) * df['pos']

    # 账户净值
    df['net_value'] = df['cash'] + df['profit']
    # ===计算爆仓
    # 至今持仓盈亏最小值
    df.loc[df['pos'] == 1, 'price_min'] = df['low']
    df.loc[df['pos'] == -1, 'price_min'] = df['high']
    df['profit_min'] = face_value * df['contract_num'] * (df['price_min'] - df['open_pos_price']) * df['pos']

    # 账户净值最小值
    df['net_value_min'] = df['cash'] + df['profit_min']
    # 计算保证金率
    df['margin_ratio'] = df['net_value_min'] / (face_value * df['contract_num'] * df['price_min'])

    # 计算是否爆仓
    df.loc[df['margin_ratio'] <= (min_margin_ratio + c_rate), '是否爆仓'] = 1

    # ===平仓时扣除手续费
    df.loc[close_pos_condition, 'net_value'] -= df['close_pos_fee']

    # 应对偶然情况：下一根K线开盘价格价格突变，在平仓的时候爆仓。此处处理有省略，不够精确。
    df.loc[close_pos_condition & (df['net_value'] < 0), '是否爆仓'] = 1

    # ===对爆仓进行处理
    df['是否爆仓'] = df.groupby('start_time')['是否爆仓'].fillna(method='ffill')
    df.loc[df['是否爆仓'] == 1, 'net_value'] = 0

    # =====计算资金曲线
    df['equity_change'] = df['net_value'].pct_change()
    df.loc[open_pos_condition, 'equity_change'] = df.loc[open_pos_condition, 'net_value'] / initial_cash - 1  # 开仓日的收益率
    df['equity_change'].fillna(value=0, inplace=True)
    df['equity_curve'] = (1 + df['equity_change']).cumprod()

    return df

# okex交割合约（usdt本位）资金曲线
def equity_curve_for_OKEx_USDT_future_next_open_kaka(df, slippage=1 / 1000, slippage2=1 / 1000, c_rate=5 / 10000, leverage_rate=3, face_value=0.01,
                                                min_margin_ratio=1/100):

    # =====找出开仓、平仓的k线(开多）
    condition1 = df['pos'] != 0  # 当前周期不为空仓
    condition2 = df['pos'] != df['pos'].shift(1)  # 当前周期和上个周期持仓方向不一样。
    open_pos_condition = condition1 & condition2

    condition1 = df['pos'] != 0  # 当前周期不为空仓（开空）
    condition2 = df['pos'] != df['pos'].shift(-1)  # 当前周期和下个周期持仓方向不一样。
    close_pos_condition = condition1 & condition2

    # =====对每次交易进行分组
    df.loc[open_pos_condition, 'start_time'] = df['candle_begin_time']
    # print(df)
    df['start_time'].fillna(method='ffill', inplace=True)
    df.loc[df['pos'] == 0, 'start_time'] = pd.NaT

    # =====开始计算资金曲线
    initial_cash = 10000  # 初始资金，默认为10000元
    # ===在开仓时
    # df.loc[open_pos_condition, 'open_price'] = df['open_price'].shift()
    df.loc[open_pos_condition, 'open_price'] = df['open_price'].shift()
    # print(df)
    # 在open_pos_condition的K线，以开盘价计算买入合约的数量。（当资金量大的时候，可以用5分钟均价）
    df.loc[open_pos_condition, 'contract_num'] = initial_cash * leverage_rate / (face_value * df['open_price'])
    # print(df)
    df['contract_num'] = np.floor(df['contract_num'])  # 对合约张数向下取整
    # 开仓价格：理论开盘价加上相应滑点
    df.loc[open_pos_condition, 'open_pos_price'] = df['open_price'] * (1 + slippage * df['pos'])
    # 开仓之后剩余的钱，扣除手续费
    df['cash'] = initial_cash - df['open_pos_price'] * face_value * df['contract_num'] * c_rate  # 即保证金

    # ===开仓之后每根K线结束时
    # 买入之后cash，contract_num，open_pos_price不再发生变动
    for _ in ['contract_num', 'open_pos_price', 'cash']:
        df[_].fillna(method='ffill', inplace=True)
    df.loc[df['pos'] == 0, ['contract_num', 'open_pos_price', 'cash']] = None

    # ===在平仓时
    # 平仓价格
    df.loc[close_pos_condition, 'close_pos_price'] = df['close_price'] * (1 - slippage2 * df['pos'])
    # 平仓之后剩余的钱，扣除手续费
    df.loc[close_pos_condition, 'close_pos_fee'] = df['close_pos_price'] * face_value * df['contract_num'] * c_rate
    # df.to_csv('ttt.csv', encoding='utf_8_sig')
    # ===计算利润
    # 开仓至今持仓盈亏
    df['profit'] = face_value * df['contract_num'] * (df['close'] - df['open_pos_price']) * df['pos']
    # 平仓时理论额外处理
    df.loc[close_pos_condition, 'profit'] = face_value * df['contract_num'] * (
            df['close_pos_price'] - df['open_pos_price']) * df['pos']
    # 账户净值
    df['net_value'] = df['cash'] + df['profit']

    # 去除profit为空的数据
    df['profit'].fillna(value=0, inplace=True)
    df = df[df['profit'] != 0]
    # print(df)

    # ===计算爆仓
    # 至今持仓盈亏最小值
    df.loc[df['pos'] == 1, 'price_min'] = df['close_pos_price']
    df.loc[df['pos'] == -1, 'price_min'] = df['close_pos_price']

    df['profit_min'] = face_value * df['contract_num'] * (df['price_min'] - df['open_pos_price']) * df['pos']
    # 账户净值最小值
    df['net_value_min'] = df['cash'] + df['profit_min']
    # 计算保证金率
    df['margin_ratio'] = df['net_value_min'] / (face_value * df['contract_num'] * df['price_min'])
    # 计算是否爆仓
    df.loc[df['margin_ratio'] <= (min_margin_ratio + c_rate), '是否爆仓'] = 1

    # ===平仓时扣除手续费
    df.loc[close_pos_condition, 'net_value'] -= df['close_pos_fee']
    # 应对偶然情况：下一根K线开盘价格价格突变，在平仓的时候爆仓。此处处理有省略，不够精确。
    df.loc[close_pos_condition & (df['net_value'] < 0), '是否爆仓'] = 1

    # ===对爆仓进行处理
    df['是否爆仓'] = df.groupby('start_time')['是否爆仓'].fillna(method='ffill')
    df.loc[df['是否爆仓'] == 1, 'net_value'] = 0

    # =====计算资金曲线
    df['equity_change'] = df['net_value'].pct_change()
    df.loc[open_pos_condition, 'equity_change'] = df.loc[open_pos_condition, 'net_value'] / initial_cash - 1  # 开仓日的收益率
    df['equity_change'].fillna(value=0, inplace=True)
    df['equity_curve'] = (1 + df['equity_change']).cumprod()

    return df


# 新策略的okex交割合约(usdt本位)资金曲线
def equity_curve_for_OKEx_USDT_future_next_open_kaka2(df, slippage=1 / 1000, slippage2=1 / 1000, c_rate=5 / 10000, leverage_rate=3, face_value=0.1,
                                                min_margin_ratio=1/100, stop_loss=0.03):

    stop_loss_price_list_reverse = []
    stop_loss_sum_reverse = 0
    stop_loss_price_list = []
    stop_loss_sum = 0
    auto_close = 1  # 1 周期内没有止损点 0 周期内有止损点

    # 找出所有反向操作止损点
    for index in range(len(df)):
        columnvalues = df.iloc[index, :].values
        # print(columnvalues)
        if columnvalues[18] == 0:
            # 一个周期中没有止损点时在中轨平仓
            if index > 0 and df.iloc[index-1, :].values[18] != 0 and auto_close != 0:
                df.loc[index+1, 'stop_loss_reverse'] = 1
            #初始化
            stop_loss_price_list_reverse = []
            stop_loss_sum_reverse = 0

        if columnvalues[18] == 1:
            auto_close = 1
            # 反向操作开多情况下开盘价 * (1-止损百分比) = 止损价
            if len(stop_loss_price_list_reverse) < 1:
                stop_loss_price_list_reverse.append(columnvalues[1] * (1-stop_loss))
                # 收盘价 < 止损价
                if columnvalues[4] < stop_loss_price_list_reverse[0]:
                    stop_loss_sum_reverse += 1
                    # 找到下根k线，标记止损点
                    df.loc[index+1, 'stop_loss_reverse'] = 1
                    auto_close -= 1
            else:
                if columnvalues[4] < stop_loss_price_list_reverse[0]:
                    if stop_loss_sum_reverse >= 1:
                        pass
                    else:
                        stop_loss_sum_reverse += 1
                        # 找到下根k线，标记止损点
                        df.loc[index+1, 'stop_loss_reverse'] = 1
                        auto_close -= 1

        elif columnvalues[18] == -1:
            # 反向操作开空情况下开盘价 * (1+止损百分比) = 止损价
            if len(stop_loss_price_list_reverse) < 1:
                stop_loss_price_list_reverse.append(columnvalues[1] * (1 + stop_loss))
                # 收盘价 > 止损价
                if columnvalues[4] > stop_loss_price_list_reverse[0]:
                    stop_loss_sum_reverse += 1
                    # 找到下根k线，标记止损点
                    df.loc[index + 1, 'stop_loss_reverse'] = 1
            else:
                if columnvalues[4] > stop_loss_price_list_reverse[0]:
                    if stop_loss_sum_reverse >= 1:
                        pass
                    else:
                        stop_loss_sum_reverse += 1
                        # 找到下根k线，标记止损点
                        df.loc[index + 1, 'stop_loss_reverse'] = 1

    auto_close = 1  # 1 周期内没有止损点 0 周期内有止损点
    # 找出所有正向操作止损点               
    for index in range(len(df)):
        columnvalues = df.iloc[index, :].values
        # print(columnvalues)
        if columnvalues[16] == 0:
            if index > 0 and df.iloc[index-1, :].values[16] != 0 and auto_close != 0:
                df.loc[index+1, 'stop_loss'] = 1
            # 初始化
            stop_loss_price_list = []
            stop_loss_sum = 0

        if columnvalues[16] == 1:
            auto_close = 1
            # 正向操作开多情况下开盘价 * (1-止损百分比) = 止损价
            if len(stop_loss_price_list) < 1:
                stop_loss_price_list.append(columnvalues[1] * (1 - stop_loss))
                # 收盘价 < 止损价
                if columnvalues[4] < stop_loss_price_list[0]:
                    stop_loss_sum += 1
                    # 找到下根k线，标记止损点
                    df.loc[index + 1, 'stop_loss'] = 1
                    auto_close -= 1
            else:
                if columnvalues[4] < stop_loss_price_list[0]:
                    if stop_loss_sum >= 1:
                        pass
                    else:
                        stop_loss_sum += 1
                        # 找到下根k线，标记止损点
                        df.loc[index + 1, 'stop_loss'] = 1
                        auto_close -= 1

        elif columnvalues[16] == -1:
            auto_close = 1
            # 正向操作开空情况下开盘价 * (1+止损百分比) = 止损价
            if len(stop_loss_price_list) < 1:
                stop_loss_price_list.append(columnvalues[1] * (1 + stop_loss))
                # 收盘价 > 止损价
                if columnvalues[4] > stop_loss_price_list[0]:
                    stop_loss_sum += 1
                    # 找到下根k线，标记止损点
                    df.loc[index + 1, 'stop_loss'] = 1
                    auto_close -= 1
            else:
                if columnvalues[4] > stop_loss_price_list[0]:
                    if stop_loss_sum >= 1:
                        pass
                    else:
                        stop_loss_sum += 1
                        # 找到下根k线，标记止损点
                        df.loc[index + 1, 'stop_loss'] = 1
                        auto_close -= 1
    # ==========================输出==========================
    # pd.set_option('display.max_rows', None)
    # print(df)
    # df.to_csv(path_or_buf='data.csv')

    # =====找出开仓的k线
    condition1 = df['pos'] != 0  # 当前周期不为空仓
    condition2 = df['pos'] != df['pos'].shift(1)  # 当前周期和上个周期持仓方向不一样。
    open_pos_condition = condition1 & condition2

    # =====对每次交易进行分组
    # df.loc[open_pos_condition, 'start_time'] = df['candle_begin_time']
    # df['start_time'].fillna(method='ffill', inplace=True)
    # df.loc[df['pos'] == 0, 'start_time'] = pd.NaT
    df.loc[open_pos_condition, 'stop_loss'] = 0
    df.loc[open_pos_condition, 'stop_loss_reverse'] = 0

    # 找出开仓点到止损点的区间
    condition1 = df['stop_loss'].shift(1) == 1
    df.loc[condition1, 'stop_loss'] = 3
    condition1 = df['stop_loss_reverse'].shift(1) == 1
    df.loc[condition1, 'stop_loss_reverse'] = 3
    df['stop_loss'].fillna(method='ffill', inplace=True)
    df['stop_loss_reverse'].fillna(method='ffill', inplace=True)

    # =====开始计算资金曲线
    initial_cash = 10000  # 初始资金，默认为10000元
    # ===在开仓时
    # 在open_pos_condition的K线，以开盘价计算买入合约的数量。（当资金量大的时候，可以用5分钟均价）
    df.loc[open_pos_condition, 'contract_num'] = initial_cash * leverage_rate / (face_value * df['open'])
    df['contract_num'] = np.floor(df['contract_num'])  # 对合约张数向下取整
    # 开仓价格：理论开盘价加上相应滑点
    df.loc[open_pos_condition, 'open_pos_price'] = df['open'] * (1 + slippage * df['pos'])
    # 本金扣除手续费后的钱
    df['cash'] = initial_cash - df['open_pos_price'] * face_value * df['contract_num'] * c_rate  # 即保证金

    # ===开仓之后每根K线结束时
    # 买入之后cash，contract_num，open_pos_price不再发生变动
    for _ in ['contract_num', 'open_pos_price', 'cash']:
        df[_].fillna(method='ffill', inplace=True)
    df.loc[df['pos'] == 0, ['contract_num', 'open_pos_price', 'cash']] = None

    # ===在平仓时
    # 平仓价格
    df.loc[df['stop_loss'] == 1, 'close_pos_price'] = df['open'] * (1 - slippage * df['pos'].shift(1))
    # 反向平仓时的价格
    df.loc[df['stop_loss_reverse'] == 1, 'close_pos_price_reverse'] = df['open'] * (1 - slippage * df['pos_reverse'].shift(1))
    # 平仓之后的手续费
    df.loc[df['stop_loss'] == 1, 'close_pos_fee'] = df['close_pos_price'] * face_value * df['contract_num'] * c_rate
    df.loc[df['stop_loss_reverse'] == 1, 'close_pos_fee_reverse'] = df['close_pos_price_reverse'] * face_value * df['contract_num'] * c_rate

    # ===计算利润
    # 开仓至今持仓盈亏
    df.loc[df['stop_loss'] == 0, 'profit'] = face_value * df['contract_num'] * (df['close'] - df['open_pos_price']) * df['pos']
    # # 反向至今持仓盈亏
    df.loc[df['stop_loss_reverse'] == 0, 'profit_reverse'] = face_value * df['contract_num'] * (df['close'] - df['open_pos_price']) * df['pos_reverse']

    # 平仓时理论额外处理
    df.loc[df['stop_loss'] == 1, 'profit'] = face_value * df['contract_num'] * (
            df['close_pos_price'] - df['open_pos_price']) * df['pos']

    # 反向平仓时理论额外处理
    df.loc[df['stop_loss_reverse'] == 1, 'profit_reverse'] = face_value * df['contract_num'] * (
            df['close_pos_price_reverse'] - df['open_pos_price']) * df['pos_reverse']

    # 账户净值
    df.loc[df['stop_loss'] == 0, 'net_value'] = df['cash'] + df['profit']
    # 反向操作账户净值
    df.loc[df['stop_loss_reverse'] == 0, 'net_value_reverse'] = df['cash'] + df['profit_reverse']

    # # ===平仓时扣除手续费
    df.loc[df['stop_loss'] == 1, 'net_value'] -= df['close_pos_fee']
    df.loc[df['stop_loss_reverse'] == 1, 'net_value_reverse'] -= df['close_pos_fee_reverse']

    # =====计算资金曲线
    df['equity_change'] = df['net_value'].pct_change()
    df.loc[open_pos_condition, 'equity_change'] = df.loc[open_pos_condition, 'net_value'] / initial_cash - 1  # 开仓日的收益率
    df['equity_change'].fillna(value=0, inplace=True)
    df['equity_curve'] = (1 + df['equity_change']).cumprod()

    df['equity_change_reverse'] = df['net_value_reverse'].pct_change()
    df.loc[open_pos_condition, 'equity_change_reverse'] = df.loc[open_pos_condition, 'net_value_reverse'] / initial_cash - 1  # 开仓日的收益率
    df['equity_change_reverse'].fillna(value=0, inplace=True)
    df['equity_curve_reverse'] = (1 + df['equity_change_reverse']).cumprod()


    # df.to_csv(path_or_buf='data.csv')
    # print(df.loc[:, ['candle_begin_time', 'open', 'close','pos_reverse','stop_loss_price','sell','pingcang', 'pingcang_price','pingcang_fee','profit_reverse','net_value_reverse']])

    # print(df)
    return df